feat: allow Tensor.store API to receive .var as value #120
Conversation
|
After this patch and #111, the |
|
I think a better approach would be to simplify the variable by looking it up in the environment, as the resulting store (with the variable) doesn't correspond to anything the HW can do. At the KLR Core level, we should not have anything that does not have a corresponding ISA (or BIR) representation. However, the problem here is more fundamental. The core issue is that is non-sensical because The above "statement form" would trace with no issues. Of course, there are other operators, like |
This patch allows Tensor.store API (which is connected to `nki.language.store`) to accept a more generic `Core.Value` type. The motivation is tracing of interop/test/examples/matmul.py, specifically the `nki_matmul_basic_` function. After apply leanprover#111, tracing the Python function was raising the following error message: ``` error: line 44: nl.store(result[i_out_p, i_out_f], value=result_sbuf) ^-- expecting tensor access ``` It is because its `value` keyword argument was having the following expression: ``` KLR.Trace.Term.expr (KLR.Core.Expr.value (KLR.Core.Value.var "5")) (KLR.Trace.TermType.obj `object) ``` which could not be converted to Access through the FromNKI typeclass. The "5" temporary variable was emerging from the right hand side of the definition of `result_sbuf`: ``` result_sbuf = nl.copy(result_psum, dtype=result.dtype) ``` To convert the value of "5", it seems we need to get the generated trace and find assignment to "5" because: ``` def RValue : Term -> Trace Term ... | .expr e@(.call ..) ty => do let v := (<- genName).toString add_stmt (.assign v e) return .expr (.value $ .var v) ty ``` the `add_stmt` is just adding a Core statement to `State.body`. Skimming through `State.body` and finding this assignment to "5" didn't seem something we wanted to do inside Tensor.store, so instead I slightly chose a conservative approach and simply removed the shape checker. But any other reasonable option is still fine with me.
aqjune-aws
changed the title
This patch allows Tensor.store API (which is connected to
nki.language.store) to accept a more genericCore.Valuetype.The motivation is tracing of interop/test/examples/matmul.py, specifically the
nki_matmul_basic_function.After apply #111, tracing the Python function was raising the following error message:
It is because its
valuekeyword argument was having the following expression:which could not be converted to Access through the FromNKI typeclass.
The "5" temporary variable was emerging from the right hand side of the definition of
result_sbuf:To convert the value of "5", it seems we need to get the generated trace and find assignment to "5" because:
the
add_stmtis just adding a Core statement toState.body.Skimming through
State.bodyand finding this assignment to "5" didn't seem something we wanted to do inside Tensor.store, so instead I slightly chose a conservative approach and simply removed the shape checker.But any other reasonable option is still fine with me.